#include <cmath>
#include <cstdio>
#include <cstdlib>

#include "DeformationState.h"
#include "FEMModel.h"
#include "TopologicalOperation.h"
#include "Tetrahedron.h"
#include "TetrahedronState.h"
//#include "LinearCGState.h"
#include "BigVector.h"
#include "Types.h"
#include "Scenario.h"
#include "DeformationState.h"
#include "Node.h"
#include "Statistics.h"
#include "Triangle.h"
#include "Mechanics.h"
#include "ModelIO.h"

const Real MEGA_REDUX = 1e6;

DeformationState::DeformationState (FiniteElementModel*fem)
{
	need_reinit_b_ = true;
	external_force_len_sq_ =0.0;
	incremental_b_ = false;

	fem_ = fem;

	char const * s = GetStringSetting ("elasticity");

	force_function_ = 0;  
	dforce_function_ = 0;
	dforce_precompute_ = 0;
	incremental_dforce_ = 0;
	energy_func_= 0;

	if (!strcmp (s, "nonlinear"))
	{
 		force_function_ =  &PicinbonoElasticForce;
 		dforce_function_ = &PicinbonoForceDerivative;
 		energy_func_ = &PicinbonoElasticEnergy;
	}
	else if (!strcmp (s, "picinbono2"))
	{
 		force_function_ =  &Picinbono2ElasticForce;
 		dforce_function_ = &Picinbono2ForceDerivative;
	}
	else if (!strcmp (s, "linear"))
	{
  		force_function_ =  &LinearElasticForce;
  		energy_func_ = &LinearElasticEnergy;
	}
	else if (!strcmp (s, "neohooke"))
	{
 		force_function_ = &NeoHookeanElasticForce ;
 		dforce_function_ = &NeoHookeanElasticForceDerivative;
 		energy_func_ = &NeoHookeanElasticEnergy;      
#if 0
		dforce_precompute_ = &NeoHookeanPrecompute;
		incremental_dforce_ = &NeoHookeanIncrementalDForce;
#endif
	}
	else if (!strcmp (s, "veronda"))
	{
 		force_function_ = &VerondaElasticForce;
 		dforce_function_ = &VerondaElasticForceDerivative;
 		energy_func_ = &VerondaElasticEnergy;            
	}
	else if (!strcmp (s, "consistent-veronda"))
	{
 		force_function_ = &ConsistentVerondaElasticForce;
 		dforce_function_ = &ConsistentVerondaElasticForceDerivative;
 		energy_func_ = &ConsistentVerondaElasticEnergy;                  
	}
	else if (!strcmp (s, "springs"))
	{
		force_function_ = 0;
		dforce_function_ = 0;
		energy_func_ = 0;
	}
	else
	{
		fprintf (stderr, "Unknown elasticity type");
		exit (2);
	}

	CompletizeNodes ();
}


/*
Sync the vectors with the node structure.  
*/
void DeformationState::UpdateVectorSizes ()
{
	/*
	To be sure that the equations are well defined, we remove unused
	nodes.

	UGH: triggers assertion fails in the cut-selector; we remove nodes
	from under the cut selector.  */
	RemoveUnusedNodes (fem_);

	ValidateNodeIncidence (fem_);
	ValidateCutRegularity (fem_);
	CompletizeNodes ();

	tet_states_.Clear ();

	LinkArray<Tetrahedron> tets = fem_->TetrahedronArray ();
	for (int i = tets.Size () ;i--;)
	{
		tets[i]->Precompute (fem_);
		tet_states_.Push (tets[i]->state_);
		tets[i]->state_->Precompute (tets[i]);
	}

	tet_states_.DefaultSort ();	// hopefully speeds up mem access ? 
}

LinkArray< Array<Real> >
DeformationState::VectorVariables ()
{
	Array<Real> *vecs[] = {
		&displacement_, &external_force_, &delta_external_force_, &stats_residual_, 0
	};

	LinkArray< Array<Real> > vvs;
	for (int i =0; vecs[i]; i++)
		vvs.Push (vecs[i]);

	return vvs;
}

void DeformationState::CompletizeNodes ()
{
	int n =fem_->NodeCount () *3;
	LinkArray<Array<Real> > vecs =  VectorVariables ();
	for (int i= 0; i < vecs.Size (); i++)
	{
		CompletizeBigVector (vecs[i],n);
	}

	while (fixed_arr_.Size () < n)
	{
		fixed_arr_.Push (false);
		outer_fix_.Push (false);
		normal_fix_.Push (false);
	}
}


void DeformationState::RemoveNodes (Array<int> const * nod_indices)
{
	LinkArray<Array<Real> > vecs =  VectorVariables ();
	for (int i= 0; i < vecs.Size (); i++)
	{
		BigVectorRemoveEntries (vecs[i],nod_indices);
	}

	BigVectorRemoveEntries (&fixed_arr_, nod_indices);
	BigVectorRemoveEntries (&outer_fix_,nod_indices);
	BigVectorRemoveEntries (&normal_fix_, nod_indices);
}


void DeformationState::ScheduleReinit ()
{
	need_reinit_b_ = true;
}

/*
update FORCE and DELTA_FORCE.
*/
void DeformationState::UpdateForces ()
{
	Real * ef = external_force_.UnsafeAccessArray();
	Real *df = delta_external_force_.UnsafeAccessArray();
	bool *fixp = (bool*) fixed_arr_.AccessArray ();
	int n = delta_external_force_.Size ();

	BigVectorPartialNullify  (df, df, fixp,n);
	BigVectorAdd (ef, df, ef, n);

	external_force_len_sq_ = BigVectorIP (ef,ef,n);

	/*
	should reset DELTA_EXTERNAL_FORCE_ after this.
	*/
}


void DeformationState::FixNode (int idx, bool fix)
{
	for (int j=0; j < 3; j++)
	{
		if (fixed_arr_[idx + j] && !fix)
			delta_external_force_[idx + j] *= -1;
		fixed_arr_[idx  + j] = fix;
	}
}

void DeformationState::AddForce (int idx, Real df)
{
	delta_external_force_[idx] += df;
}

void DeformationState::DuplicateVariable (int dst, int src)
{
	displacement_[dst] = displacement_[src];
}



void DeformationState::ScheduleMoveNode (int n , Real r)
{

}

int DeformationState::Dimension () const
{
	return external_force_.Size ();
}

void DeformationState::DoSimulationBody()
{
}

void DeformationState::SimulationBody ()
{
	/*
	Check if vector sizes are up to date
	*/
	if (need_reinit_b_)
	{
		need_reinit_b_ = 0;
		UpdateVectorSizes ();
	}

	UpdateForces ();

	DoSimulationBody ();


	BigVectorNullify (delta_external_force_.UnsafeAccessArray (), Dimension ());



	/*

	*/

#if 0
	UpdateGeometryStuff (fem_);
#endif
}


#if 0
void UpdateGeometryStuff (FiniteElementModel*fem)
{
	void * coll_data = fem->collision_data_;
	Triangle *tri = [...] ;

	Node * node0 = tri->GetNode (0);

	Vector3 loc_vec =  node0->Location (true); 
	Real const * displacements = fem->pstate_->displacement_.AccessArray ();
	Real const * locations = fem->node_locations_.AccessArray ();

	double loc[3];
	for (int i =0; i < 3; i++)
	{
		loc[i] = displacements[3 * node0->Number () + i] + locations[3 * node0->Number () + i];
	}
}

#endif

void DeformationState::InitSimulation ()
{
	CompletizeNodes ();

	LinkArray<Tetrahedron> ts (fem_->TetrahedronArray ());
	for(int i = ts.Size (); i--;)
		ts[i]->Precompute(fem_);

	ScheduleReinit ();
}

DeformationState::~DeformationState ()
{

}

void DeformationState::CheckComponents (ComponentMap const*)
{

}

int force_apply_count = 0;
long long flop_count = 0LL ;
void
DeformationState::ApplyForceFunction (Real *dest,  Real const * src) const
{
	BigVectorNullify (dest, Dimension ());
	if (force_function_)
	{
		TetrahedronState ** tetra_start = tet_states_.AccessArray ();
		TetrahedronState ** tetra_end = tetra_start + tet_states_.Size ();

		for (TetrahedronState ** tp = tetra_start; tp < tetra_end; tp++)
		{
			flop_count += (*tp)->ElasticForce (force_function_, dest, src);
		}
	}
}

/*
Calculate the derivative K(displ) * dir
*/
void DeformationState::ApplyForceDerivative (Real *deriv, Real *residual_vec, Real const*displ, Real const *dir) const
{
	int n = Dimension ();
	BigVectorNullify (deriv, n);
	if (residual_vec)
		BigVectorNullify (residual_vec,n);

	if (dforce_function_)
	{
		TetrahedronState ** tetra_start = tet_states_.AccessArray ();
		TetrahedronState ** tetra_end = tetra_start + tet_states_.Size ();

		if (incremental_b_ && incremental_dforce_)
			for (TetrahedronState ** tp = tetra_start; tp < tetra_end; tp++)
			{
				flop_count += (*tp)->IncrementalDForce (incremental_dforce_, deriv, dir);
			}
		else for (TetrahedronState ** tp = tetra_start; tp < tetra_end; tp++)
		{
			flop_count +=
				(*tp)->ElasticForceDerivative (dforce_function_, deriv, residual_vec, displ, dir);
		}
	}
	else
	{
		ApplyForceFunction (deriv, dir);
		if (residual_vec)
			ApplyForceFunction (residual_vec, displ);
	}
}

void DeformationState::ApplyForceDerivativeAndResidual (Real *deriv, Real *residual_vec, Real const*displ, Real const *dir) const
{
	/*
	elasticity
	*/

	ApplyForceDerivative (deriv, residual_vec, displ, dir);
	int n = Dimension();	

	bool const *fixp  = fixed_arr_.AccessArray ();  

	if (residual_vec)
	{
		/*
		and external_force
		*/

		Real const *external_force = external_force_.AccessArray ();

		BigVectorAdd (residual_vec, residual_vec, external_force, n);
		BigVectorPartialNullify (residual_vec, residual_vec, fixp, n);
	}

	/*
	for consistency: we clear the deriv of fixed nodes as well. 
	*/
	BigVectorPartialNullify (deriv, deriv, fixp, n);
}

/*
should move into static-deformation-state? 
*/
Real DeformationState::ComputeResidualForce (Real *residual_dest, Real const *displacement) const
{
	/*
	Residu is  (stiffness  * displ )
	*/
	ApplyForceFunction (residual_dest, displacement); 

	// + external_force
	bool *fixp  = (bool*) fixed_arr_.AccessArray ();
	Real *force = (Real*) external_force_.AccessArray ();
	int n =  Dimension ();

	BigVectorAdd (residual_dest, residual_dest, force, n);
	BigVectorPartialNullify (residual_dest, residual_dest, fixp, n);

	flop_count += n;

	return BigVectorIP (residual_dest,residual_dest,n);
}


void DeformationState::FlopStats () const
{
	long long remember = flop_count ;

	printf ("MFLOPS: %f",  (flop_count / MEGA_REDUX) ); 

	Real residual_len
		=sqrt (ComputeResidualForce ((Real*)stats_residual_.AccessArray (),
		displacement_.AccessArray ()));

	Real relative_res = residual_len / sqrt(external_force_len_sq_);

	printf(" REL_FORCE_LEN: %g SOL_DIFFERENCE: %g", relative_res, DeformationStateDistance (this));

	Real e  = VirtualEnergy ();

	printf( " ENERGY: %f\n", e);

	flop_count = remember;
}

void DeformationState::UpdateIncrementalStep ()
{
	if (!dforce_precompute_)
		return ;

	TetrahedronState ** tetra_start = tet_states_.AccessArray ();
	TetrahedronState ** tetra_end = tetra_start + tet_states_.Size ();

	for (TetrahedronState ** tp = tetra_start; tp < tetra_end; tp++)
	{
		flop_count += (*tp)->PrecomputeForDisplacement (dforce_precompute_,
			displacement_.AccessArray ());
	}
}


void DeformationState::ValidateNumericalSanity ()
{
	Real const *d = displacement_.AccessArray ();
	int n = displacement_.Size();
	bool b = BigVectorSane (d, n);

	b = b && (BigVectorLength (d,n) < 1e15);
	if (!b)
	{
		fprintf(stderr, "Numerical La-La land reached. Kaboom.\n");
		exit (2);
	}
}

Real DeformationState::VirtualEnergy ()const
{
	Real e = 0.0;
	if (energy_func_)
	{

		Real const *disp = 	   displacement_.AccessArray ();

		e=  - BigVectorIP (external_force_.AccessArray (),
			disp,
			Dimension());
		TetrahedronState ** tetra_start = tet_states_.AccessArray ();
		TetrahedronState ** tetra_end = tetra_start + tet_states_.Size ();

		for (TetrahedronState ** tp = tetra_start; tp < tetra_end; tp++)
		{
			e += (*tp)->ElasticEnergy (energy_func_, disp);
		}
	}
	return e;
}
